use std::fmt;
use std::fmt::Write;

use crate::lexical::{Node, Token};

impl fmt::Display for Node {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let mut ctx = FormatContext::new();
        self.print_code(f, &mut ctx)?;
        Ok(())
    }
}

fn write_indent(f: &mut fmt::Formatter, level: usize) -> fmt::Result {
    let indent = "    ".repeat(level);
    write!(f, "{}", indent)?;
    Ok(())
}

struct FormatContext {
    indent: usize,
    last_token: Option<Token>,
    last_keyword: Option<Token>,
    last_indent: usize,
}

impl FormatContext {
    pub fn new() -> Self {
        FormatContext {
            indent: 0,
            last_token: None,
            last_keyword: None,
            last_indent: 0,
        }
    }
}

impl Node {
    fn print_children(
        &self,
        f: &mut fmt::Formatter,
        ctx: &mut FormatContext,
    ) -> fmt::Result {
        for child in &self.children {
            if let Some(node) = child {
                node.borrow().print_code(f, ctx)?;
            }
        }
        Ok(())
    }

    fn print_code(
        &self,
        f: &mut fmt::Formatter,
        ctx: &mut FormatContext,
    ) -> fmt::Result {
        let mut keyword = None;
        let mut o = String::new();
        let res =
            match &self.token {
                Token::KeyWord(s) => {
                    if s == "else" {
                        write_indent(f, ctx.indent)?;
                    }
                    write!(o, "{}", s)?;
                    keyword = Some(self.token.clone());
                    Ok(())
                }
                Token::Plus(s) => {
                  if self.has_parent(Token::UnaryOp, 1) {
                    write!(o, " {}", s)?;
                  } else {
                    write!(o, " {} ", s)?;
                  }
                  Ok(())
                }
                Token::Minus(s) => {
                  if self.has_parent(Token::UnaryOp, 1) {
                    write!(o, " {}", s)?;
                  } else {
                    write!(o, " {} ", s)?;
                  }
                  Ok(())
                }
                Token::Mul(s) => write!(o, " {} ", s),
                Token::Div(s) => write!(o, " {} ", s),
                Token::Mod(s) => write!(o, " {} ", s),
                Token::Assign(s) => write!(o, " {} ", s),
                Token::Eq(s) => write!(o, " {} ", s),
                Token::Neq(s) => write!(o, " {} ", s),
                Token::Lt(s) => write!(o, " {} ", s),
                Token::Gt(s) => write!(o, " {} ", s),
                Token::Le(s) => write!(o, " {} ", s),
                Token::Ge(s) => write!(o, " {} ", s),
                Token::Not(s) => write!(o, "{}", s),
                Token::And(s) => write!(o, " {} ", s),
                Token::Or(s) => write!(o, " {} ", s),
                Token::LParen(s) => write!(o, "{}", s),
                Token::RParen(s) => write!(o, "{}", s),
                Token::LBrace(_) => {
                    if let Some(t) = &ctx.last_token {
                        if *t == Token::RParen(")".to_string()) {
                            write!(o, " ")?;
                        }
                    }
                    // DEBUG
                    // if let Some(t) = &ctx.last_token {
                    //     write!(o, "{}", t)?;
                    // }
                    if self.has_parent(Token::Decl, 100) {
                      write!(o, "{{")?;
                    } else {
                      write!(o, "{{\n")?;
                    }
                    Ok(())
                }
                Token::RBrace(_) => {
                    if self.has_parent(Token::Decl, 100) {
                        write!(o, "}}")?;
                    } else {
                        if ctx.indent > 0 {
                          write_indent(f, ctx.indent - 1)?;
                        }
                        write!(o, "}}\n")?;
                    }
                    Ok(())
                }
                Token::LBrackt(s) => write!(o, "{}", s),
                Token::RBrackt(s) => write!(o, "{}", s),
                Token::Comma(s) => write!(o, "{} ", s),
                Token::Semicolon(s) => write!(o, "{}\n", s),
                Token::Ident(s) => write!(o, "{}", s),

                Token::Root => self.print_children(f, ctx),
                Token::CompUnit => {
                    for (i, child) in self.children.iter().enumerate() {
                        if let Some(node) = child {
                            node.borrow().print_code(f, ctx)?;
                            // Write empty line.
                            if i < self.children.len() - 1 {
                                writeln!(f)?;
                            }
                        }
                    }
                    Ok(())
                }
                Token::Decl => {
                    write_indent(f, ctx.indent)?;
                    self.print_children(f, ctx)?;
                    Ok(())
                }
                Token::BType(t) => {
                    write!(o, "{} ", t)?;
                    Ok(())
                }
                Token::FuncFParams => self.print_children(f, ctx),
                Token::FuncFParam => self.print_children(f, ctx),
                Token::FuncDef => {
                    self.print_children(f, ctx)?;
                    writeln!(f)
                }
                Token::FuncType(ty) => {
                    write!(o, "{} ", ty)?;
                    Ok(())
                }
                Token::Block => {
                    for child in &self.children {
                        if let Some(node) = child {
                            ctx.indent += 1;
                            node.borrow().print_code(f, ctx)?;
                            ctx.indent -= 1;
                        }
                    }
                    Ok(())
                }
                Token::IfStmt => self.print_children(f, ctx),
                Token::WhileStmt => self.print_children(f, ctx),
                Token::Stmt => {
                    // DEBUG
                    // if let Some(t) = &ctx.last_token {
                    //     write!(o, "{}", t)?;
                    // }
                    if ctx.last_keyword.as_ref().map_or(false, |t| {
                        *t == Token::KeyWord("else".to_string())
                    }) {
                        self.print_children(f, ctx)?;
                    } else {
                        if !(ctx.last_token.as_ref().map_or(false, |t| {
                            *t == Token::RParen(")".to_string())
                        })) {
                            write_indent(f, ctx.indent)?;
                        }
                        self.print_children(f, ctx)?;
                    }
                    Ok(())
                }
                Token::Exp => self.print_children(f, ctx),
                Token::Number(num) => write!(o, "{}", num),
                Token::UnaryOp => self.print_children(f, ctx),
                Token::ConstExp => self.print_children(f, ctx),
                _ => {
                    write!(o, "{}", self.token)?;
                    Ok(())
                }
            };
        if ctx
            .last_keyword
            .as_ref()
            .map_or(false, |t| *t == Token::KeyWord("else".to_string()))
        {
            if keyword
                .as_ref()
                .map_or(false, |t| *t == Token::KeyWord("if".to_string()))
            {
                write!(f, " ")?;
            } else {
                writeln!(f)?;
                write_indent(f, ctx.indent + 1)?;
            }
        }
        write!(f, "{}", o)?;
        ctx.last_token = Some(self.token.clone());
        ctx.last_keyword = keyword;
        ctx.last_indent = ctx.indent;
        res
    }
}
